#ifndef _HOOK_SYSCALL_INTERNAL_H
#define _HOOK_SYSCALL_INTERNAL_H

#include <asm/pgtable.h>  // {clear,set}_pte_bit(), set_pte()
#include <linux/errno.h>
#include <linux/sched.h>
#include <linux/unistd.h>
#include <linux/version.h>

#ifdef __aarch64__
#include <asm/tlbflush.h>   // flush_tlb_kernel_range()
#include <linux/mm.h>       // struct mm_struct, apply_to_page_range()
#include <linux/vmalloc.h>  // vm_unmap_aliases()
static struct mm_struct* init_mm_ptr = 0;
#endif

#include "lookup_addr.h"
#include "syscall_hook_proto.h"

#define HOOK(syscall_name) static sys_##syscall_name##_func_t syscall_name = 0;
#define SYSCALL_HOOK_INIT(old, new)                                       \
    {                                                                     \
        .syscall = (unsigned long*)&old, .hook_func = (unsigned long)new, \
        __NR_##old                                                        \
    }
/**
 * @brief hook结构体，用于原始调用，自定义逻辑，系统调用编号表示。
 *
 */
struct syscall_hook_list {
    unsigned long* syscall;
    unsigned long  hook_func;
    int            NR;
};
typedef struct syscall_hook_list syscall_hook_list_t;

/**
 * @brief hook系统调用所需上下文
 *
 */
struct syscall_hook_ctx {
    unsigned long sys_call_table_addr;
    unsigned int  level;
    pte_t*        pte;
};
typedef struct syscall_hook_ctx syscall_hook_ctx_t;

#ifdef __x86_64__

static int sh_set_page_rw(syscall_hook_ctx_t* ctx)
{
    set_pte(ctx->pte, pte_mkwrite(*(ctx->pte)));
    return 0;
}

static int sh_set_page_ro(syscall_hook_ctx_t* ctx)
{
    set_pte(ctx->pte, pte_wrprotect(*(ctx->pte)));
    return 0;
}

#elif __aarch64__

#if 0
static pte_t *lookup_address_arm64(unsigned long addr)
{
	pgd_t *pgd;
	pud_t *pud;
	pmd_t *pmd;
	pte_t *pte;
	struct mm_struct *sh_init_mm;

	sh_init_mm = (struct mm_struct*)sh_lookup_addr("init_mm");
	if(0 == sh_init_mm) return NULL;
	
	pgd = pgd_offset(sh_init_mm, addr);
	if (pgd_none(*pgd)) return NULL;

	pud = pud_offset(pgd, addr);
	if (pud_none(*pud)) return NULL;

	pmd = pmd_offset(pud, addr);
	if (pmd_none(*pmd)) return NULL;

	pte = pte_offset_kernel(pmd, addr);
	if (!pte_valid(*pte))
		return NULL;

	return pte;
}
#endif

/********** HELPERS **********/
// quote:
// https://stackoverflow.com/questions/61247838/cannot-use-set-memory-rw-in-linux-kernel-on-arm64#

struct page_change_data {
    pgprot_t set_mask;
    pgprot_t clear_mask;
};

static int change_page_range(pte_t* ptep, pgtable_t token, unsigned long addr,
                             void* data)
{
    struct page_change_data* cdata = data;
    pte_t                    pte   = READ_ONCE(*ptep);

    pte = clear_pte_bit(pte, cdata->clear_mask);
    pte = set_pte_bit(pte, cdata->set_mask);

    set_pte(ptep, pte);
    return 0;
}

static int __change_memory_common(unsigned long start, unsigned long size,
                                  pgprot_t set_mask, pgprot_t clear_mask)
{
    struct page_change_data data;
    int                     ret;

    data.set_mask   = set_mask;
    data.clear_mask = clear_mask;

    ret =
        apply_to_page_range(init_mm_ptr, start, size, change_page_range, &data);

    flush_tlb_kernel_range(start, start + size);
    return ret;
}

static int set_page_rw(unsigned long addr)
{
    vm_unmap_aliases();
    return __change_memory_common(addr, PAGE_SIZE, __pgprot(PTE_WRITE),
                                  __pgprot(PTE_RDONLY));
}

static int set_page_ro(unsigned long addr)
{
    vm_unmap_aliases();
    return __change_memory_common(addr, PAGE_SIZE, __pgprot(PTE_RDONLY),
                                  __pgprot(PTE_WRITE));
}

static int sh_set_page_rw(syscall_hook_ctx_t* ctx)
{
    unsigned long addr = ctx->sys_call_table_addr;
    addr               = addr & PAGE_MASK;
    return set_page_rw(addr);
}

static int sh_set_page_ro(syscall_hook_ctx_t* ctx)
{
    unsigned long addr = ctx->sys_call_table_addr;
    addr               = addr & PAGE_MASK;
    return set_page_ro(addr);
}

#endif

/**
 * @brief hook 初始化
 *
 * @param ctx hook上下文结构体
 * @param list hook结构体数组
 * @param len 数组长度
 * @return int
 */
static int syscall_hook_init(syscall_hook_ctx_t* ctx, syscall_hook_list_t* list,
                             int len)
{
    int i;
#if LINUX_VERSION_CODE < KERNEL_VERSION(2, 6, 12)
    printk(
        "the kernel version less than 2.6.12 is not suport because "
        "`kallsyms_lookup_name` is not export\n");
    return -EBADRQC;
#endif

    ctx->sys_call_table_addr = sh_lookup_addr("sys_call_table");
    if (ctx->sys_call_table_addr == 0) {
        printk("- unable to locate sys_call_table\n");
        return -EFAULT;
    }
#ifdef __x86_64__
    ctx->pte =
        lookup_address((unsigned long)ctx->sys_call_table_addr, &ctx->level);
    if (NULL == ctx->pte) {
        printk("page not found.\n");
        return -EFAULT;
    }
#elif __aarch64__
    init_mm_ptr = (struct mm_struct*)sh_lookup_addr("init_mm");
    if (init_mm_ptr == 0) {
        printk("init_mm not found.\n");
        return -EFAULT;
    }
#endif

    sh_set_page_rw(ctx);
    for (i = 0; i < len; i++) {
        *list[i].syscall =
            ((unsigned long*)(ctx->sys_call_table_addr))[list[i].NR];
        ((unsigned long*)(ctx->sys_call_table_addr))[list[i].NR] =
            list[i].hook_func;
    }
    sh_set_page_ro(ctx);

    return 0;
}

/**
 * @brief
 *
 * @param ctx hook清理
 * @param list
 * @param len
 * @return int
 */
static int syscall_hook_clean(syscall_hook_ctx_t*  ctx,
                              syscall_hook_list_t* list, int len)
{
    int i;
    sh_set_page_rw(ctx);
    for (i = 0; i < len; i++) {
        ((unsigned long*)(ctx->sys_call_table_addr))[list[i].NR] =
            *list[i].syscall;
    }
    sh_set_page_ro(ctx);
    return 0;
}

#endif /* _HOOK_SYSCALL_INTERNAL_H */
